package com.jumper.mydemo.service.impl;

import org.springframework.stereotype.Service;

import com.jumper.mydemo.service.TongYiService;

import reactor.core.publisher.Flux;

import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.image.ImageClient;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;

import java.util.Map;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.springframework.ai.chat.ChatClient;
import org.springframework.ai.chat.StreamingChatClient;


@Service
public class TongYiServiceImpl implements TongYiService {

    private final static Logger logger = LogManager.getLogger(TongYiServiceImpl.class);

    private final ChatClient chatClient;

    private final StreamingChatClient streamingChatClient;

    private final ImageClient imageClient;

    @Value("classpath:/prompts/zxw.st")
    private Resource zxwResource;

    public TongYiServiceImpl(ChatClient chatClient, StreamingChatClient streamingChatClient, ImageClient imageClient) {
        this.chatClient = chatClient;
        this.streamingChatClient = streamingChatClient;
        this.imageClient = imageClient;
    }

    @Override
    public String completion(String message) {
        //PromptTemplate template = new PromptTemplate(zxwResource);
        Prompt prompt = new Prompt(message);
        logger.info("user input: {}", message);
        return chatClient.call(prompt).getResult().getOutput().getContent();
    }

    @Override
    public Map<String, String> streamCompletion(String message) {
        StringBuilder fullContent = new StringBuilder();
        streamingChatClient.stream(new Prompt(message))
            .flatMap(chatResponse -> Flux.fromIterable(chatResponse.getResults()))
            .map(content -> content.getOutput().getContent())
            .doOnNext(fullContent::append)
            .last()
            .map(lastContent -> Map.of(message, fullContent.toString()))
            .block();

        return Map.of(message, fullContent.toString());
    }

    @Override
    public ImageResponse genImage(String imgPrompt) {
        return imageClient.call(new ImagePrompt(imgPrompt));
    }
    
}
